package com.ruoyi.computing.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjUtil;
import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson2.JSONObject;
import com.aliyun.oss.model.OSSObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.ruoyi.common.enums.cms.GenerateFileType;
import com.ruoyi.common.enums.cms.OriginalFileType;
import com.ruoyi.common.utils.oss.OssUtil;
import com.ruoyi.common.utils.video.VideoFrameExtractor;
import com.ruoyi.computing.entity.domain.ComputingModel;
import com.ruoyi.common.enums.cms.ModelStatusEnum;
import com.ruoyi.computing.entity.domain.ComputingModelGenerateFile;
import com.ruoyi.computing.entity.domain.ComputingModelOriginalFile;
import com.ruoyi.computing.mapper.ComputingModelGenerateFileMapper;
import com.ruoyi.computing.mapper.ComputingModelOriginalFileMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import com.ruoyi.computing.mapper.ComputingModelMapper;
import com.ruoyi.computing.service.ComputingModelService;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.net.InetAddress;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 模型Service业务层处理
 */
@Service
@Slf4j
public class ComputingModelServiceImpl implements ComputingModelService {

    @Autowired
    private ComputingModelMapper modelMapper;
    @Autowired
    private ComputingModelOriginalFileMapper modelOriginalFileMapper;
    @Autowired
    private ComputingModelGenerateFileMapper modelGenerateFileMapper;


    @Value("${ruoyi.env}")
    private String env;
    @Value("${gaussian-splatting.path}")
    private String gsPath;
    @Value("${oss.modelPath}")
    private String modelPath;

    @Scheduled(cron = "0 0/3 * * * ?")
    @Override
    public void generateModel() throws Exception {
        if (!env.startsWith("gpu")) {
            return;
        }
        InetAddress localhost = InetAddress.getLocalHost();
        String ip = localhost.getHostAddress();
        log.info("当前IP: " + ip);

        LambdaQueryWrapper<ComputingModel> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.in(ComputingModel::getStatus, ModelStatusEnum.QUEUING.getStatus(), ModelStatusEnum.GENERATING.getStatus());
        queryWrapper.eq(ComputingModel::getServerIp, ip);
        queryWrapper.eq(ComputingModel::getDeleted, 0);
        queryWrapper.orderByAsc(ComputingModel::getId);

        List<ComputingModel> computingModels = modelMapper.selectList(queryWrapper);

        if (CollUtil.isNotEmpty(computingModels)) {
            Map<Integer, List<ComputingModel>> collect = computingModels.stream().collect(Collectors.groupingBy(ComputingModel::getStatus));

            //生成中模型
            List<ComputingModel> generatingList = collect.get(ModelStatusEnum.GENERATING.getStatus());
            if (CollUtil.isNotEmpty(generatingList)) {
                log.info("模型生成中...");
            } else {
                //排队中模型
                List<ComputingModel> queuingList = collect.get(ModelStatusEnum.QUEUING.getStatus());

                ComputingModel computingModel = queuingList.get(0);
                queuingProcess(computingModel);
            }
        }
    }

    private void queuingProcess(ComputingModel computingModel) throws Exception {
        log.info("模型生成start");
        //更新状态为生成中
        computingModel.setStatus(ModelStatusEnum.GENERATING.getStatus());
        modelMapper.updateById(computingModel);

        try {
            // 1.处理文件夹
            handleFolder();
            // 2.素材文件下载到本地
            handleOriginalFile(computingModel);
            // 3.调用python脚本接口，开始训练模型
            startTrain();
            // 4.上传模型文件
            uploadModel(computingModel);

        } catch (Exception e) {
            computingModel.setStatus(ModelStatusEnum.FAIL.getStatus());
            modelMapper.updateById(computingModel);
            log.error("模型生成失败", e);
            throw e;
        }

        //更新状态为成功
        computingModel.setStatus(ModelStatusEnum.SUCCESS.getStatus());
        modelMapper.updateById(computingModel);
    }

    /**
     * 处理文件夹
     */
    private void handleFolder() throws Exception {
        //删除文件夹
        Path path = Paths.get(gsPath + "/data");
        try {
            Files.walk(path)
                    .sorted(Comparator.reverseOrder())
                    .map(Path::toFile)
                    .forEach(File::delete);
            log.info("文件夹删除成功: " + path);
        } catch (Exception e) {
            log.error("文件夹删除失败: " + path, e);
            throw e;
        }

        //新建文件夹
        try {
            // 创建多级文件夹
            Files.createDirectories(path);
            log.info("文件夹创建成功: " + path);

            // 创建子文件夹
            Path subFolder = path.resolve("input");
            Files.createDirectory(subFolder);
            System.out.println("子文件夹创建成功: " + subFolder);

        } catch (IOException e) {
            log.info("文件夹创建失败", e);
            throw e;
        }
    }

    /**
     * 素材文件下载到本地
     */
    private void handleOriginalFile(ComputingModel computingModel) throws Exception {
        //查询素材
        LambdaQueryWrapper<ComputingModelOriginalFile> queryWrapper = new LambdaQueryWrapper<>();
        queryWrapper.eq(ComputingModelOriginalFile::getModelId, computingModel.getId());
        List<ComputingModelOriginalFile> computingModelOriginalFiles = modelOriginalFileMapper.selectList(queryWrapper);

        if (CollUtil.isEmpty(computingModelOriginalFiles)) {
            log.error("无素材");
            throw new Exception("无素材");
        }

        Integer originalType = computingModel.getOriginalType();

        //oss下载文件 存到本地
        try {
            if (OriginalFileType.IMAGES.getType().equals(originalType)) {
                //图片批量下载
                for (ComputingModelOriginalFile file : computingModelOriginalFiles) {
                    String result = file.getObjectName().substring(file.getObjectName().lastIndexOf("/"));
                    OssUtil.downloadLocal(file.getObjectName(), gsPath + "/data/input" + result);
                }
                computingModel.setOriginalNum(computingModelOriginalFiles.size());
            } else if (OriginalFileType.VIDEO.getType().equals(originalType)) {
                ComputingModelOriginalFile file = computingModelOriginalFiles.get(0);
                //视频只下载一个
                String result = file.getObjectName().substring(file.getObjectName().lastIndexOf("/"));
                OssUtil.downloadLocal(file.getObjectName(), gsPath + "/data/input" + result);
                //视频抽帧
                int frameInterval = 10;
                if (null != computingModel.getFrameInterval()) {
                    frameInterval = computingModel.getFrameInterval();
                }
                int originalNum = VideoFrameExtractor.extractFrames(gsPath + "/data/input" + result, gsPath + "/data/input", frameInterval);
                computingModel.setOriginalNum(originalNum);
            }
            modelMapper.updateById(computingModel);
        } catch (Exception e) {
            log.error("下载文件失败", e);
            throw e;
        }
    }

    /**
     * 开始训练模型
     */
    private void startTrain() throws Exception {
        // 1.图片素材预处理
        pyHttp("http://localhost:5000/py/api/convert");
        log.info("图片素材预处理完成");
        // 2.预处理（不处理会报错）
        pyHttp("http://localhost:5000/py/api/fix_error");
        log.info("报错预处理完成");
        // 3.训练模型
        pyHttp("http://localhost:5000/py/api/train");
        log.info("训练模型完成");
        // 4.格式转换（ply转splat）
        pyHttp("http://localhost:5000/py/api/splat");
        log.info("格式转换完成");
    }

    private void pyHttp(String url) throws Exception {
        String result = HttpUtil.post(url, new HashMap<>());
        JSONObject jsonObject = JSONObject.parseObject(result);
        if (ObjUtil.isEmpty(jsonObject) || !(0 == jsonObject.getIntValue("return_code"))) {
            throw new Exception("调用python接口失败");
        }
    }

    /**
     * 上传模型文件
     */
    private void uploadModel(ComputingModel computingModel) throws Exception {
        // 7000 ply
        uploadPly(computingModel.getId(), "/iteration_7000", GenerateFileType.PLY_7000);

        // 30000 ply
        uploadPly(computingModel.getId(), "/iteration_30000", GenerateFileType.PLY_30000);

        // 7000 splat
        uploadPly(computingModel.getId(), "/iteration_7000", GenerateFileType.SPLAT_7000);

        // 30000 splat
        Long modelSize = uploadPly(computingModel.getId(), "/iteration_30000", GenerateFileType.SPLAT_30000);
        computingModel.setModelSize(modelSize);
        modelMapper.updateById(computingModel);
    }

    /**
     * 上传模型文件
     */
    private Long uploadPly(Long modelId, String folder, GenerateFileType type) throws Exception {
        Date now = new Date();

        UUID uuid = UUID.randomUUID();
        String objectName = modelPath + uuid + "." + type.getSuffix();

        OSSObject ossObject = OssUtil.upload(objectName, gsPath + "/data/output/point_cloud" + folder + "/" + type.getName());

        ComputingModelGenerateFile modelGenerateFile = new ComputingModelGenerateFile();
        modelGenerateFile.setModelId(modelId);
        modelGenerateFile.setUrl("https://tcsy-oss.oss-cn-guangzhou.aliyuncs.com/" + objectName);
        modelGenerateFile.setObjectName(objectName);
        modelGenerateFile.setType(type.getType());
        modelGenerateFile.setSize(ossObject.getObjectMetadata().getContentLength());
        modelGenerateFile.setCreator(0L);
        modelGenerateFile.setCreateTime(now);
        modelGenerateFileMapper.insert(modelGenerateFile);

        return ossObject.getObjectMetadata().getContentLength();
    }


}
